from typing import Any
from warnings import warn
from collections import Iterable
import numpy as np
import torch
from torch import Tensor
from .typing import *

__all__ = ['is_listy','is_tuple', 'is_dict',
    'tensor', 'normalize', 'denormalize', 'image_to_tensor', 'mask_to_tensor', 'cat']

def is_listy(x:Any)->bool: return isinstance(x, (tuple,list))
def is_tuple(x:Any)->bool: return isinstance(x, tuple)
def is_dict(x:Any)->bool: return isinstance(x, dict)

def listify(p:OptListOrItem=None, q:OptListOrItem=None):
    "Make `p` listy and the same length as `q`."
    if p is None: p=[]
    elif isinstance(p, str):          p = [p]
    elif not isinstance(p, Iterable): p = [p]
    #Rank 0 tensors in PyTorch are Iterable but don't have a length.
    else:
        try: a = len(p)
        except: p = [p]
    n = q if type(q)==int else len(p) if q is None else len(q)
    if len(p)==1: p = p * n
    assert len(p)==n, f'List len mismatch ({len(p)} vs {n})'
    return list(p)

def cat(tensors:Tensors, dim=0):
    """
    Efficient version of torch.cat that avoids a copy if there is only a single element in a list
    """
    assert isinstance(tensors, (list, tuple))
    if len(tensors) == 1:
        return tensors[0]
    return torch.cat(tensors, dim)


def tensor(x:Any, *rest)->Tensor:
    "Like `torch.as_tensor`, but handle lists too, and can pass multiple vector elements directly."
    if len(rest): x = (x,)+rest
    # XXX: Pytorch bug in dataloader using num_workers>0; TODO: create repro and report
    if is_listy(x) and len(x)==0: return tensor(0)
    res = torch.tensor(x)
    if res.dtype is torch.int32:
        warn('Tensor is int32: upgrading to int64; for better performance use int64 input')
        return res.long()
    return res

def normalize(x:TensorImage, mean:Tensors,std:Tensors)->TensorImage:
    mean, std = tensor(mean), tensor(std)
    "Normalize `x` with `mean` and `std`."
    return (x-mean[...,None,None]) / std[...,None,None]

def denormalize(x:TensorImage, mean:Tensors,std:Tensors, do_x:bool=True)->TensorImage:
    "Denormalize `x` with `mean` and `std`."
    mean, std = tensor(mean), tensor(std)
    return x.cpu().float()*std[...,None,None] + mean[...,None,None] if do_x else x.cpu()

def image_to_tensor(im, stats=None):
    tensor = torch.from_numpy(np.moveaxis(im / (255. if im.dtype == np.uint8 else 1), -1, 0).astype(np.float32))
    if stats is not None:
        return normalize(tensor, *stats)
    return tensor


def mask_to_tensor(mask):
    # mask = np.expand_dims(mask / (255. if mask.dtype == np.uint8 else 1), 0).astype(np.float32)
    mask = (mask / (255. if mask.dtype == np.uint8 else 1)).astype(np.float32)
    return torch.from_numpy(mask).long()